/*----------------------------------------------------------------------------*/
/*  CP2K: A general program to perform molecular dynamics simulations         */
/*  Copyright 2000-2025 CP2K developers group <https://cp2k.org>              */
/*                                                                            */
/*  SPDX-License-Identifier: BSD-3-Clause                                     */
/*----------------------------------------------------------------------------*/
#include "cp_mpi.h"

#include <assert.h>
#include <omp.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#if defined(__parallel)
/*******************************************************************************
 * \brief Check given MPI status and upon failure abort with a nice message.
 * \author Ole Schuett
 ******************************************************************************/
#define CHECK(CMD)                                                             \
  do {                                                                         \
    const int error = (CMD);                                                   \
    if (MPI_SUCCESS != error) {                                                \
      fprintf(stderr, "MPI error #%i in %s:%i\n", error, __FILE__, __LINE__);  \
      MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE);                                 \
    }                                                                          \
  } while (0)
#endif

/*******************************************************************************
 * \brief Wrapper around MPI_Init.
 * \author Ole Schuett
 ******************************************************************************/
void cp_mpi_init(int *argc, char ***argv) {
#if defined(__parallel)
  CHECK(MPI_Init(argc, argv));
#else
  (void)argc; // mark used
  (void)argv;
#endif
}

/*******************************************************************************
 * \brief Wrapper around MPI_Finalize.
 * \author Ole Schuett
 ******************************************************************************/
void cp_mpi_finalize(void) {
#if defined(__parallel)
  CHECK(MPI_Finalize());
#endif
}

/*******************************************************************************
 * \brief Returns MPI_COMM_WORLD.
 * \author Ole Schuett
 ******************************************************************************/
cp_mpi_comm_t cp_mpi_get_comm_world(void) {
#if defined(__parallel)
  return MPI_COMM_WORLD;
#else
  return -1;
#endif
}

/*******************************************************************************
 * \brief Wrapper around MPI_Comm_f2c.
 * \author Ole Schuett
 ******************************************************************************/
cp_mpi_comm_t cp_mpi_comm_f2c(const int fortran_comm) {
#if defined(__parallel)
  // Attempt to support "no MPI" if __parallel is defined (0 -> MPI_COMM_NULL).
  return 0 != fortran_comm ? MPI_Comm_f2c(fortran_comm) : MPI_COMM_NULL;
#else
  (void)fortran_comm; // mark used
  return -1;
#endif
}

/*******************************************************************************
 * \brief Wrapper around MPI_Comm_c2f.
 * \author Ole Schuett
 ******************************************************************************/
int cp_mpi_comm_c2f(const cp_mpi_comm_t comm) {
#if defined(__parallel)
  return MPI_Comm_c2f(comm);
#else
  (void)comm; // mark used
  return -1;
#endif
}

/*******************************************************************************
 * \brief Wrapper around MPI_Comm_rank.
 * \author Ole Schuett
 ******************************************************************************/
int cp_mpi_comm_rank(const cp_mpi_comm_t comm) {
#if defined(__parallel)
  int rank = 0;
  if (MPI_COMM_NULL != comm) { // !MPI_Comm_compare
    CHECK(MPI_Comm_rank(comm, &rank));
  }
  return rank;
#else
  (void)comm; // mark used
  return 0;
#endif
}

/*******************************************************************************
 * \brief Wrapper around MPI_Comm_size.
 * \author Ole Schuett
 ******************************************************************************/
int cp_mpi_comm_size(const cp_mpi_comm_t comm) {
#if defined(__parallel)
  int nranks = 1;
  if (MPI_COMM_NULL != comm) { // !MPI_Comm_compare
    CHECK(MPI_Comm_size(comm, &nranks));
  }
  return nranks;
#else
  (void)comm; // mark used
  return 1;
#endif
}

/*******************************************************************************
 * \brief Wrapper around MPI_Dims_create.
 * \author Ole Schuett
 ******************************************************************************/
void cp_mpi_dims_create(const int nnodes, const int ndims, int dims[]) {
#if defined(__parallel)
  CHECK(MPI_Dims_create(nnodes, ndims, dims));
#else
  dims[0] = nnodes;
  for (int i = 1; i < ndims; i++) {
    dims[i] = 1;
  }
#endif
}

/*******************************************************************************
 * \brief Wrapper around MPI_Cart_create.
 * \author Ole Schuett
 ******************************************************************************/
cp_mpi_comm_t cp_mpi_cart_create(const cp_mpi_comm_t comm_old, const int ndims,
                                 const int dims[], const int periods[],
                                 const int reorder) {
#if defined(__parallel)
  cp_mpi_comm_t comm_cart;
  CHECK(MPI_Cart_create(comm_old, ndims, dims, periods, reorder, &comm_cart));
  return comm_cart;
#else
  (void)comm_old; // mark used
  (void)ndims;
  (void)dims;
  (void)periods;
  (void)reorder;
  return -1;
#endif
}

/*******************************************************************************
 * \brief Wrapper around MPI_Cart_get.
 * \author Ole Schuett
 ******************************************************************************/
void cp_mpi_cart_get(const cp_mpi_comm_t comm, int maxdims, int dims[],
                     int periods[], int coords[]) {
#if defined(__parallel)
  CHECK(MPI_Cart_get(comm, maxdims, dims, periods, coords));
#else
  (void)comm; // mark used
  for (int i = 0; i < maxdims; i++) {
    dims[i] = 1;
    periods[i] = 1;
    coords[i] = 0;
  }
#endif
}

/*******************************************************************************
 * \brief Wrapper around MPI_Cart_rank.
 * \author Ole Schuett
 ******************************************************************************/
int cp_mpi_cart_rank(const cp_mpi_comm_t comm, const int coords[]) {
#if defined(__parallel)
  int rank;
  CHECK(MPI_Cart_rank(comm, coords, &rank));
  return rank;
#else
  (void)comm; // mark used
  (void)coords;
  return 0;
#endif
}

/*******************************************************************************
 * \brief Wrapper around MPI_Cart_sub.
 * \author Ole Schuett
 ******************************************************************************/
cp_mpi_comm_t cp_mpi_cart_sub(const cp_mpi_comm_t comm,
                              const int remain_dims[]) {
#if defined(__parallel)
  cp_mpi_comm_t newcomm;
  CHECK(MPI_Cart_sub(comm, remain_dims, &newcomm));
  return newcomm;
#else
  (void)comm; // mark used
  (void)remain_dims;
  return -1;
#endif
}

/*******************************************************************************
 * \brief Wrapper around MPI_Comm_free.
 * \author Ole Schuett
 ******************************************************************************/
void cp_mpi_comm_free(cp_mpi_comm_t *comm) {
#if defined(__parallel)
  CHECK(MPI_Comm_free(comm));
#else
  (void)comm; // mark used
#endif
}

/*******************************************************************************
 * \brief Wrapper around MPI_Comm_compare.
 * \author Ole Schuett
 ******************************************************************************/
bool cp_mpi_comms_are_similar(const cp_mpi_comm_t comm1,
                              const cp_mpi_comm_t comm2) {
#if defined(__parallel)
  int res;
  CHECK(MPI_Comm_compare(comm1, comm2, &res));
  return res == MPI_IDENT || res == MPI_CONGRUENT || res == MPI_SIMILAR;
#else
  (void)comm1; // mark used
  (void)comm2;
  return true;
#endif
}

/*******************************************************************************
 * \brief Wrapper around MPI_Allreduce for op MPI_MAX and datatype MPI_INT.
 * \author Ole Schuett
 ******************************************************************************/
void cp_mpi_max_int(int *values, const int count, const cp_mpi_comm_t comm) {
#if defined(__parallel)
  if (MPI_COMM_NULL != comm) { // !MPI_Comm_compare
    int value = 0;
    void *recvbuf =
        (1 < count ? cp_mpi_alloc_mem(count * sizeof(int)) : &value);
    CHECK(MPI_Allreduce(values, recvbuf, count, MPI_INT, MPI_MAX, comm));
    memcpy(values, recvbuf, count * sizeof(int));
    if (1 < count) {
      cp_mpi_free_mem(recvbuf);
    }
  }
#else
  (void)comm; // mark used
  (void)values;
  (void)count;
#endif
}

/*******************************************************************************
 * \brief Wrapper around MPI_Allreduce for op MPI_MAX and datatype MPI_UINT64_T.
 * \author Ole Schuett
 ******************************************************************************/
void cp_mpi_max_uint64(uint64_t *values, const int count,
                       const cp_mpi_comm_t comm) {
#if defined(__parallel)
  if (MPI_COMM_NULL != comm) { // !MPI_Comm_compare
    uint64_t value = 0;
    void *recvbuf =
        (1 < count ? cp_mpi_alloc_mem(count * sizeof(uint64_t)) : &value);
    CHECK(MPI_Allreduce(values, recvbuf, count, MPI_UINT64_T, MPI_MAX, comm));
    memcpy(values, recvbuf, count * sizeof(uint64_t));
    if (1 < count) {
      cp_mpi_free_mem(recvbuf);
    }
  }
#else
  (void)comm; // mark used
  (void)values;
  (void)count;
#endif
}

/*******************************************************************************
 * \brief Wrapper around MPI_Allreduce for op MPI_MAX and datatype MPI_DOUBLE.
 * \author Ole Schuett
 ******************************************************************************/
void cp_mpi_max_double(double *values, const int count,
                       const cp_mpi_comm_t comm) {
#if defined(__parallel)
  if (MPI_COMM_NULL != comm) { // !MPI_Comm_compare
    double value = 0;
    void *recvbuf =
        (1 < count ? cp_mpi_alloc_mem(count * sizeof(double)) : &value);
    CHECK(MPI_Allreduce(values, recvbuf, count, MPI_DOUBLE, MPI_MAX, comm));
    memcpy(values, recvbuf, count * sizeof(double));
    if (1 < count) {
      cp_mpi_free_mem(recvbuf);
    }
  }
#else
  (void)comm; // mark used
  (void)values;
  (void)count;
#endif
}

/*******************************************************************************
 * \brief Wrapper around MPI_Allreduce for op MPI_SUM and datatype MPI_INT.
 * \author Ole Schuett
 ******************************************************************************/
void cp_mpi_sum_int(int *values, const int count, const cp_mpi_comm_t comm) {
#if defined(__parallel)
  if (MPI_COMM_NULL != comm) { // !MPI_Comm_compare
    int value = 0;
    void *recvbuf =
        (1 < count ? cp_mpi_alloc_mem(count * sizeof(int)) : &value);
    CHECK(MPI_Allreduce(values, recvbuf, count, MPI_INT, MPI_SUM, comm));
    memcpy(values, recvbuf, count * sizeof(int));
    if (1 < count) {
      cp_mpi_free_mem(recvbuf);
    }
  }
#else
  (void)comm; // mark used
  (void)values;
  (void)count;
#endif
}

/*******************************************************************************
 * \brief Wrapper around MPI_Allreduce for op MPI_SUM and datatype MPI_LONG.
 * \author Hans Pabst
 ******************************************************************************/
void cp_mpi_sum_long(long *values, const int count, const cp_mpi_comm_t comm) {
#if defined(__parallel)
  if (MPI_COMM_NULL != comm) { // !MPI_Comm_compare
    long value = 0;
    void *recvbuf =
        (1 < count ? cp_mpi_alloc_mem(count * sizeof(long)) : &value);
    CHECK(MPI_Allreduce(values, recvbuf, count, MPI_LONG, MPI_SUM, comm));
    memcpy(values, recvbuf, count * sizeof(long));
    if (1 < count) {
      cp_mpi_free_mem(recvbuf);
    }
  }
#else
  (void)comm; // mark used
  (void)values;
  (void)count;
#endif
}

/*******************************************************************************
 * \brief Wrapper around MPI_Allreduce for op MPI_SUM and datatype MPI_INT64_T.
 * \author Ole Schuett
 ******************************************************************************/
void cp_mpi_sum_int64(int64_t *values, const int count,
                      const cp_mpi_comm_t comm) {
#if defined(__parallel)
  if (MPI_COMM_NULL != comm) { // !MPI_Comm_compare
    int64_t value = 0;
    void *recvbuf =
        (1 < count ? cp_mpi_alloc_mem(count * sizeof(int64_t)) : &value);
    CHECK(MPI_Allreduce(values, recvbuf, count, MPI_INT64_T, MPI_SUM, comm));
    memcpy(values, recvbuf, count * sizeof(int64_t));
    if (1 < count) {
      cp_mpi_free_mem(recvbuf);
    }
  }
#else
  (void)comm; // mark used
  (void)values;
  (void)count;
#endif
}

/*******************************************************************************
 * \brief Wrapper around MPI_Allreduce for op MPI_SUM and datatype MPI_DOUBLE.
 * \author Ole Schuett
 ******************************************************************************/
void cp_mpi_sum_double(double *values, const int count,
                       const cp_mpi_comm_t comm) {
#if defined(__parallel)
  if (MPI_COMM_NULL != comm) { // !MPI_Comm_compare
    double value = 0;
    void *recvbuf =
        (1 < count ? cp_mpi_alloc_mem(count * sizeof(double)) : &value);
    CHECK(MPI_Allreduce(values, recvbuf, count, MPI_DOUBLE, MPI_SUM, comm));
    memcpy(values, recvbuf, count * sizeof(double));
    if (1 < count) {
      cp_mpi_free_mem(recvbuf);
    }
  }
#else
  (void)comm; // mark used
  (void)values;
  (void)count;
#endif
}

/*******************************************************************************
 * \brief Wrapper around MPI_Sendrecv for datatype MPI_BYTE.
 * \author Ole Schuett
 ******************************************************************************/
int cp_mpi_sendrecv_byte(const void *sendbuf, const int sendcount,
                         const int dest, const int sendtag, void *recvbuf,
                         const int recvcount, const int source,
                         const int recvtag, const cp_mpi_comm_t comm) {
#if defined(__parallel)
  MPI_Status status;
  CHECK(MPI_Sendrecv(sendbuf, sendcount, MPI_BYTE, dest, sendtag, recvbuf,
                     recvcount, MPI_BYTE, source, recvtag, comm, &status));
  int count_received = 0;
  CHECK(MPI_Get_count(&status, MPI_BYTE, &count_received));
  return count_received;
#else
  (void)sendbuf; // mark used
  (void)sendcount;
  (void)dest;
  (void)sendtag;
  (void)recvbuf;
  (void)recvcount;
  (void)source;
  (void)recvtag;
  (void)comm;
  fprintf(stderr, "Error: cp_mpi_sendrecv_byte not available without MPI\n");
  abort();
#endif
}

/*******************************************************************************
 * \brief Wrapper around MPI_Sendrecv for datatype MPI_DOUBLE.
 * \author Ole Schuett
 ******************************************************************************/
int cp_mpi_sendrecv_double(const double *sendbuf, const int sendcount,
                           const int dest, const int sendtag, double *recvbuf,
                           const int recvcount, const int source,
                           const int recvtag, const cp_mpi_comm_t comm) {
#if defined(__parallel)
  MPI_Status status;
  CHECK(MPI_Sendrecv(sendbuf, sendcount, MPI_DOUBLE, dest, sendtag, recvbuf,
                     recvcount, MPI_DOUBLE, source, recvtag, comm, &status));
  int count_received;
  CHECK(MPI_Get_count(&status, MPI_DOUBLE, &count_received));
  return count_received;
#else
  (void)sendbuf; // mark used
  (void)sendcount;
  (void)dest;
  (void)sendtag;
  (void)recvbuf;
  (void)recvcount;
  (void)source;
  (void)recvtag;
  (void)comm;
  fprintf(stderr, "Error: cp_mpi_sendrecv_double not available without MPI\n");
  abort();
#endif
}

/*******************************************************************************
 * \brief Wrapper around MPI_Alltoall for datatype MPI_INT.
 * \author Ole Schuett
 ******************************************************************************/
void cp_mpi_alltoall_int(const int *sendbuf, const int sendcount, int *recvbuf,
                         const int recvcount, const cp_mpi_comm_t comm) {
#if defined(__parallel)
  CHECK(MPI_Alltoall(sendbuf, sendcount, MPI_INT, recvbuf, recvcount, MPI_INT,
                     comm));
#else
  (void)comm; // mark used
  assert(sendcount == recvcount);
  memcpy(recvbuf, sendbuf, sendcount * sizeof(int));
#endif
}

/*******************************************************************************
 * \brief Wrapper around MPI_Alltoallv for datatype MPI_BYTE.
 * \author Ole Schuett
 ******************************************************************************/
void cp_mpi_alltoallv_byte(const void *sendbuf, const int *sendcounts,
                           const int *sdispls, void *recvbuf,
                           const int *recvcounts, const int *rdispls,
                           const cp_mpi_comm_t comm) {
#if defined(__parallel)
  CHECK(MPI_Alltoallv(sendbuf, sendcounts, sdispls, MPI_BYTE, recvbuf,
                      recvcounts, rdispls, MPI_BYTE, comm));
#else
  (void)comm; // mark used
  assert(sendcounts[0] == recvcounts[0]);
  assert(sdispls[0] == 0 && rdispls[0] == 0);
  memcpy(recvbuf, sendbuf, sendcounts[0]);
#endif
}

/*******************************************************************************
 * \brief Wrapper around MPI_Alltoallv for datatype MPI_DOUBLE.
 * \author Ole Schuett
 ******************************************************************************/
void cp_mpi_alltoallv_double(const double *sendbuf, const int *sendcounts,
                             const int *sdispls, double *recvbuf,
                             const int *recvcounts, const int *rdispls,
                             const cp_mpi_comm_t comm) {
#if defined(__parallel)
  CHECK(MPI_Alltoallv(sendbuf, sendcounts, sdispls, MPI_DOUBLE, recvbuf,
                      recvcounts, rdispls, MPI_DOUBLE, comm));
#else
  (void)comm; // mark used
  assert(sendcounts[0] == recvcounts[0]);
  assert(sdispls[0] == 0 && rdispls[0] == 0);
  memcpy(recvbuf, sendbuf, sendcounts[0] * sizeof(double));
#endif
}

/*******************************************************************************
 * \brief Wrapper around MPI_Alloc_mem.
 * \author Hans Pabst
 ******************************************************************************/
void *cp_mpi_alloc_mem(size_t size) {
  void *result = NULL;
#if DBM_ALLOC_MPI && defined(__parallel)
  CHECK(MPI_Alloc_mem((MPI_Aint)size, MPI_INFO_NULL, &result));
#elif DBM_ALLOC_OPENMP && (201811 /*v5.0*/ <= _OPENMP)
  result = omp_alloc(size, omp_null_allocator);
#else
  result = malloc(size);
#endif
  return result;
}

/*******************************************************************************
 * \brief Wrapper around MPI_Free_mem.
 * \author Hans Pabst
 ******************************************************************************/
void cp_mpi_free_mem(void *mem) {
#if DBM_ALLOC_MPI && defined(__parallel)
  CHECK(MPI_Free_mem(mem));
#elif DBM_ALLOC_OPENMP && (201811 /*v5.0*/ <= _OPENMP)
  omp_free(mem, omp_null_allocator);
#else
  free(mem);
#endif
}

// EOF
